from django.db import migrations


def set_compress(apps, schema_editor):
    if schema_editor.connection.vendor != 'postgresql':
        return

    SimpleTable = apps.get_model('main', 'SimpleTable')
    schema_editor.execute(f'''ALTER TABLE {SimpleTable._meta.db_table} 
    SET (
        timescaledb.compress, 
        timescaledb.compress_orderby = 'time DESC', 
        timescaledb.compress_segmentby = 'entity_id'
    )''')
    schema_editor.execute(f"SELECT add_compression_policy('{SimpleTable._meta.db_table}', INTERVAL '1 min')")


def remove_compress(apps, schema_editor):
    if schema_editor.connection.vendor != 'postgresql':
        return

    SimpleTable = apps.get_model('main', 'SimpleTable')
    sql = f"SELECT chunk_name  FROM timescaledb_information.chunks WHERE hypertable_name = '{SimpleTable._meta.db_table}' AND is_compressed"
    with schema_editor.connection.cursor() as cursor:
        cursor.execute(sql)
        chunk_names = cursor.fetchall()

    for chunk_name in chunk_names:
        schema_editor.execute(f"SELECT decompress_chunk('_timescaledb_internal.{chunk_name[0]}');")

    schema_editor.execute(f"SELECT remove_compression_policy('{SimpleTable._meta.db_table}')")


class Migration(migrations.Migration):

    dependencies = [
        ('main', '0001_initial'),
    ]

    operations = [
        migrations.RunPython(set_compress, remove_compress),
    ]
